"""
目前做了水平镜像、亮度、平移。进度条很好，比yolo5的进度条更简洁
未作旋转、垂直镜像、超框处理

原图和生成后的图都会存到 AUG_IMG_DIR
原图和生成后的图的 yolov5 格式的标签（txt）都会存到 AUG_LAB_DIR，xml 文件不保存
"""
import os
import shutil
import xml.etree.ElementTree as ET

import imgaug as ia
import numpy as np
from imgaug import augmenters as iaa
from PIL import Image
from tqdm import tqdm

from getLabel import CLASS, saveTxt, convert, getAnnotBoxLoc

ia.seed(1)


def convertBoxV2(box, size=(200, 200)):
    """将锚框绝对大小转换为相对于图像的相对大小
    box：锚框 (xmin,ymin,xmax,ymax) 注意此顺序和 getLabel 中不一致
    size：图像大小 (w,h)
    """
    dw = 1.0 / size[0]
    dh = 1.0 / size[1]
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return str(x) + " " + str(y) + " " + str(w) + " " + str(h)


def read_xml_annotation(root, image_id):
    """读取xml文件并使用ElementTree对xml文件进行解析，找到每个object的坐标值。"""
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall("object"):  # 找到root节点下的所有country节点
        bndbox = object.find("bndbox")  # 子节点下节点rank的值

        xmin = int(bndbox.find("xmin").text)
        xmax = int(bndbox.find("xmax").text)
        ymin = int(bndbox.find("ymin").text)
        ymax = int(bndbox.find("ymax").text)
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([xmin, ymin, xmax, ymax])
        # print(bndboxlist)

    bndbox = root.find("object").find("bndbox")
    return bndboxlist


def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
    in_file = open(os.path.join(root, str(image_id) + ".xml"))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    # 修改增强后的xml文件中的filename
    elem = tree.find("filename")
    elem.text = str(id) + ".jpg"
    xmlroot = tree.getroot()
    # 修改增强后的xml文件中的path
    elem = tree.find("path")
    if elem != None:
        elem.text = saveroot + str(id) + ".jpg"

    index = 0
    for object in xmlroot.findall("object"):  # 找到root节点下的所有country节点
        bndbox = object.find("bndbox")  # 子节点下节点rank的值

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find("xmin")
        xmin.text = str(new_xmin)
        ymin = bndbox.find("ymin")
        ymin.text = str(new_ymin)
        xmax = bndbox.find("xmax")
        xmax.text = str(new_xmax)
        ymax = bndbox.find("ymax")
        ymax.text = str(new_ymax)

        index = index + 1


def mkdir(path):
    if not os.path.exists(path):  # 判断路径是否存在
        os.makedirs(path)
        print(path + " 创建成功")
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + " 目录已存在")
        return False


if __name__ == "__main__":

    IMG_DIR = "./NEU-DET/IMAGES/"  # 根据实际情况修改路径
    XML_DIR = "./NEU-DET/ANNOTATIONS"  # 根据实际情况修改路径

    # AUG_XML_DIR = "./AUG/Annotations/"  # 存储增强后的XML文件夹路径
    AUG_IMG_DIR = (
        r"/run/media/kearney/a/CAU/42course/毕设/datasets/NEU-DET/images"  # 存储增强后的影像文件夹路径
    )
    AUG_LAB_DIR = r"/run/media/kearney/a/CAU/42course/毕设/datasets/NEU-DET/labels"
    mkdir(AUG_IMG_DIR)
    mkdir(AUG_LAB_DIR)

    AUGLOOP = 10  # 每张影像增强的数量

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    # 影像增强
    seq = iaa.Sequential(
        [
            iaa.Invert(0.5),  # 在50％的图像中反转所有像素
            iaa.Fliplr(0.5),  # 镜像 水平翻转图像
            iaa.Multiply(
                (0.5, 1.5)
            ),  # 将每个图像乘以0.5到1.5之间的随机值,使图像更暗或更亮, doesn't affect BBs
            iaa.GaussianBlur(sigma=(0, 3.0)),  # 使用0到3.0的sigma模糊图像
            iaa.Affine(
                translate_px={"x": 15, "y": 15},
                scale=(0.8, 0.95),
            ),  # translate by 40/60px on x/y axis, 图像缩放 50-70%, affects BBs
        ]
    )

    for name in tqdm(os.listdir(XML_DIR), desc="Processing"):

        bndbox = read_xml_annotation(XML_DIR, name)
        print(XML_DIR, name)
        # 保存原xml文件
        # shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)  # 复制文件

        saveTxt(
            name[0:-4],
            getAnnotBoxLoc(os.path.join(XML_DIR, name)),
            AUG_LAB_DIR,
        )  # 将原 xml 文件的标签也提取出来保存为 txt

        og_img = Image.open(IMG_DIR + "/" + name[:-4] + ".jpg")

        og_img.convert("RGB").save(
            os.path.join(AUG_IMG_DIR, (name[:-4] + ".jpg")), "JPEG"
        )  # 保存原图
        og_xml = open(os.path.join(XML_DIR, name))
        tree = ET.parse(og_xml)
        # 修改增强后的xml文件中的filename
        elem = tree.find("filename")
        elem.text = name[:-4] + ".jpg"
        # tree.write(os.path.join(AUG_XML_DIR, name)) # 保存原 xml 文件

        className = tree.find("object").findtext("name")  # 获取 xml 分类中的名称
        classId = CLASS.index(className)  # 分类对于数字序号
        print(className, classId)

        for epoch in range(AUGLOOP):
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            # 读取图片
            img = Image.open(os.path.join(IMG_DIR, name[:-4] + ".jpg"))
            # sp = img.size
            img = np.asarray(img)
            # bndbox 坐标增强
            for i in range(len(bndbox)):
                bbs = ia.BoundingBoxesOnImage(
                    [
                        ia.BoundingBox(
                            x1=bndbox[i][0],
                            y1=bndbox[i][1],
                            x2=bndbox[i][2],
                            y2=bndbox[i][3],
                        ),
                    ],
                    shape=img.shape,
                )

                bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                boxes_img_aug_list.append(bbs_aug)

                # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                if n_x1 == 1 and n_x1 == n_x2:
                    n_x2 += 1
                if n_y1 == 1 and n_y2 == n_y1:
                    n_y2 += 1
                if n_x1 >= n_x2 or n_y1 >= n_y2:
                    print("error", name)
                new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
            # 存储变化后的图片
            image_aug = seq_det.augment_images([img])[0]
            path = os.path.join(
                AUG_IMG_DIR, str(str(name[:-4]) + "_" + str(epoch)) + ".jpg"
            )
            image_auged = bbs.draw_on_image(image_aug, size=0)
            Image.fromarray(image_auged).convert("RGB").save(path)  # 保存图像


            ObjBndBoxSet = ""  # 锚框序列转换结果
            for box in new_bndbox_list:
                ObjBndBoxSet += str(classId) + " " + convertBoxV2(box) + "\n"

            saveTxt(
                str(name[:-4]) + "_" + str(epoch), ObjBndBoxSet, fileptah=AUG_LAB_DIR
            )  # 保存 yolov5 格式的 txt 文件到对应目录

            # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg')
            new_bndbox_list = []
